# -*- coding: utf=8 -*-

from os import stat, times

import numpy as np
import math
import matplotlib.pyplot as plt
import time


class robot_env(object):

    def __init__(self):
        self.max_speed = 4.0
        self.min_speed = -4.5
        self.max_yawrate = 40.0
        self.accel = 1
        self.max_dyawrate = 40.0 * math.pi/180.0
        self.v_reso = 0.01
        self.yawrate_reso = 0.1
        self.dt = 0.1
        self.predict_time =3.0
        self.to_goal_cost_gain = 0.15
        self.speed_cost_gain = 1.0
        self.obstacle_cost_gain =1.0
        self.robot_radius = 0.2
        self.goal =np.random.randint(0,10,2)
        self.ob = np.random.randint(0,10,(20,2))
        self.state = np.array([0.0,0.0,math.pi/8.0,0.0,0.0])
        # traj_temp = np.array(self.state)
    def run(self,u):
        # u = u if u<self.max_speed else self.max_speed
        # u = u if u>self.min_speed else self.min_speed
        traj_temp = np.array()
        for i in range(int(self.predict_time/self.dt)):
            state_temp= self.motion(self.state,u)
            traj_temp = np.vstack((traj_temp,state_temp))
        reward,done = self.cost(traj_temp)
        return reward,state_temp,traj_temp,done

    def run_dwa(self,u,state):
        # u = u if u<self.max_speed else self.max_speed
        # u = u if u>self.min_speed else self.min_speed
        traj_temp = np.array(state)
        state_temp = state.copy()
        for i in range(int(self.predict_time/self.dt)):
            state_temp= self.motion(state_temp,u)
            traj_temp = np.vstack((traj_temp,state_temp))
        reward,done = self.cost(traj_temp)
        return reward,state_temp,traj_temp,done

    def reset(self):
        self.state = np.array([0.0,0.0,math.pi/8.0,0.0,0.0])
        self.goal = np.random.randint(0,10,2)
        self.ob = np.random.randint(0,10,(20,2))

    def show(self,traj):
        x,y= [],[]
        xb,yb = [],[]
        for i in traj:
            x.append(i[0])
            y.append(i[1])

        for i in self.ob:
            xb.append(i[0])
            yb.append(i[1])
        fig = plt.figure()

        ax = fig.gca()
        figure = ax.plot(x,y)
        plt.scatter(xb,yb)
        plt.scatter(self.goal[0],self.goal[1],c=10)
        plt.scatter(self.goal[0],self.goal[1],c=10)
        plt.show()
    

    def cost(self,traj):
        done = False
        to_goal_cost = self.to_goal_cost_gain*self.calc_to_goal_cost(traj)
        speed_cost = self.speed_cost_gain*(self.max_speed - traj[-1][3])*0
        ob_cost = self.obstacle_cost_gain*self.calc_obstacle_cost(traj)*0
        final_cost = to_goal_cost+speed_cost+ob_cost
        dist_to_goal = math.sqrt((traj[-1,0]-self.goal[0])**2+(traj[-1][1]-self.goal[1])**2)
        if dist_to_goal <= self.robot_radius:
            done = True
            # final_cost+=100
        return final_cost,done
    ## 距离目标距离作为cost,距离越大代价越大
    def calc_to_goal_cost(self,traj):
        
        dx = self.goal[0] -traj[-1][0]
        dy = self.goal[1] -traj[-1][1]
        cost = math.sqrt(dx**2 + dy**2)*100
        return cost

    def calc_obstacle_cost(self,traj):
        skip_n = 2
        minr = float("inf")
        for ii in range(0,len(traj[:,1]),skip_n):
            for i in range(len(self.ob[:,0])):
                ox=self.ob[i,0]
                oy=self.ob[i,1]
                dx = traj[ii,0]-ox
                dy = traj[ii,1]-oy
                r = math.sqrt(dx**2+dy**2)
                if r < self.robot_radius:
                    return float("Inf")
                if minr >= r:
                    minr = r
        return minr


    def motion(self,x,u):
        x[0]+= u[0] * math.cos(x[2]) * self.dt
        x[1]+= u[0] * math.sin(x[2]) * self.dt
        x[2]+= u[1]*self.dt
        x[3] = u[0]
        x[4] = u[1]
        return x



    def calc_dynamic_window(self,x):
        x = x.copy()
        speed_x_max = x[3] + self.accel*self.dt 
        speed_x_min = x[3] - self.accel*self.dt        
        speed_y_max = x[4] + self.max_dyawrate*self.dt
        speed_y_min = x[4] - self.max_dyawrate*self.dt
        
        if speed_x_max > self.max_speed:
            speed_x_max = self.max_speed
        if speed_x_min < self.min_speed:
            speed_x_min = self.min_speed
        if speed_y_max > self.max_yawrate:
            speed_y_max = self.max_yawrate
        if speed_y_min> self.max_yawrate:
            speed_y_min = self.max_yawrate 
        Vd = [speed_x_min,speed_x_max,speed_y_min,speed_y_max]
        return Vd

    ## DWA算法实现
    def dwa(self,state):
        dw = self.calc_dynamic_window(state)
        u,traj = self.calc_dwa_output(state,dw)
        return u,traj
    
    def calc_dwa_output(self,state,dw):
        state_init =state.copy()
        min_cost = float("inf")
        best_u = [0.0,0.0]
        best_traj = np.array([state])
        ## 遍历所有的线速度和角速度
        ## 找到代价最小的[v,y]
        for v in np.arange(dw[0],dw[1],self.v_reso):
            for y in np.arange(dw[2],dw[3],self.yawrate_reso): 
                reward = None
                reward,state_temp,traj_temp,_ = self.run_dwa([v,y],state_init)
                if min_cost >= reward:
                    min_cost = reward
                    best_traj = traj_temp
                    best_u = [v,y] 
        return best_u,best_traj


    def setstartandgoal(self,startgoal):
        self.goal = np.array(startgoal[2:])
        # self.ob = np.random.randint(0,10,(20,2))
        self.state = np.array([startgoal[0],startgoal[1],math.pi/8.0,0.0,0.0])       

    def test(self,algo="dwa",show = True):

        if algo == "dwa":
            time_mean = 0
            turn = 1
            for _ in range(turn):
                print("turn==",_)
                # self.reset()
                state = self.state
                u = np.array([0.0,0.0])
                traj = 0
                traj = np.array(state)
                time_start = time.time()
                dist_to_goal = float("inf")
                dist_to_goal_old = float("inf")
                count = 0
                while(1):
                    u,_ = self.dwa(state)
                    state= self.motion(state,u)
                    dist_to_goal = math.sqrt((state[0]-self.goal[0])**2+(state[1]-self.goal[1])**2)
                    x = dist_to_goal
                    traj = np.vstack((traj,state))
                    if dist_to_goal < self.robot_radius:
                        break
                    if abs(dist_to_goal_old - x)< 0.001:
                        count+=1
                        if count>10:
                            print(dist_to_goal)
                            # self.show(traj)
                            
                            break
                    dist_to_goal_old = dist_to_goal
                time_end = time.time()-time_start
                time_mean+= time_end/turn
                if show:
                    pass
                    # self.show(traj)
            print(time_mean)
        if algo == "PPO2":
            pass
        return traj

if __name__ == "__main__": 
    print("_______now_______start_______")
    robotEnv = robot_env()
    robotEnv.test()